import os
import linecache
import numpy
import random
import math

class LinUCB(object):
    
    ALPHA = 1
    
    def __init__(self, file, N):
        self.number_of_images = N
        self.input = os.path.dirname(__file__) + '/Input/' + file + str(N)
        self.images_presented = []
        self.number_of_features = len(self.GetImageVector(0))
        self.A = dict()
        self.b = dict()
        
    def GetImageVector(self, number):
        # Get vector from line #number from file input
        # number starts from 1
        line = linecache.getline(self.input, number+1)
        image = numpy.array((map(float, line.split())))
        return numpy.exp(image)
    
    def ChooseFirstImageSet(self, setsize):
        chosen_images = []
        while len(chosen_images) < setsize:
            image = random.randint(0, self.number_of_images-1)
            if image not in chosen_images:
                chosen_images.append(image)
        # save images as been presented
        self.previous_images = chosen_images        
        return chosen_images
    
    def LinUCB(self, images, relevance_scores, setsize):
        
        for image in images:
            if image not in self.A:
                self.A[image] = numpy.identity(self.number_of_features)
                self.b[image] = numpy.zeros(self.number_of_features)
            x = self.GetImageVector(image)
            self.A[image] += numpy.dot(x, x.T)
            self.b[image] += relevance_scores[image]*x.T
                
        p = []
        
        for image in range(self.number_of_images):
            if image not in self.A:
                self.A[image] = numpy.identity(self.number_of_features)
                self.b[image] = numpy.zeros(self.number_of_features)

            Theta = numpy.dot(numpy.linalg.inv(self.A[image]),self.b[image])
            x = self.GetImageVector(image)
            p.append(float(numpy.dot(Theta.T,x) + LinUCB.ALPHA*math.sqrt(numpy.dot(numpy.dot(x.T,numpy.linalg.inv(self.A[image])),x))))
                    
        chosen_images = []
        indeces = numpy.argsort(p)
        i = 0
        while len(chosen_images) < setsize:
            # starting from the last image
            image = indeces[len(indeces)-i-1]
            if image not in self.previous_images:
                chosen_images.append(image)
                self.images_presented.append(image)
            i += 1
        self.previous_images = chosen_images
                
        return chosen_images
            
            
            